# -*- coding: utf-8 -*-
"""
Author: chen xi (modified by Qwen)
Date: 2025/09/03
File: random_string_wake_client.py
功能：定时发送加密随机字符串 UDP 唤醒包，MSG_TYPE 随机为 1 或 3，持续 24 小时
"""

import struct
import socket
import time
import random
import string
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime

# ==================== 配置参数 ====================
SERVER_IP = "120.26.136.5"
SERVER_PORT = 50001
TIMEOUT = 5  # UDP 超时（秒）
NUM_THREADS = 1  # 线程池大小
REQUESTS_PER_BATCH = 50  # 每批次发送数量

# 随机字符串长度范围
MIN_STR_LEN = 1
MAX_STR_LEN = 64

# 加密参数
ENCRYPTION_KEY = "RSIOTGWakeupKey@"
ENCRYPTION_MODE = 256  # 如 AES-256

# 发送控制
SEND_INTERVAL = 0.2  # 每 0.2 秒发送一批（5 批/秒）
RUN_DURATION = 24 * 60 * 60  # 总运行时间：24 小时（秒）

# UDP 协议常量
START_CODE = 0xFFFFFFA1
MSG_HEADER_FORMAT = ">III"  # StartCode, MsgType, MsgLen (大端)


# ==================== 导入加密模块 ====================
try:
    from string_enc import iPN_StringEnc
except ImportError as e:
    print(f"❌ 无法导入 string_enc 模块: {e}")
    print("请确保 string_enc.py 在当前目录")
    exit(1)


# ==================== 工具函数 ====================
def get_random_msg_type():
    """随机返回消息类型：1 或 3"""
    return random.choice([1, 3])


def generate_random_string():
    """生成指定长度范围内的随机字符串"""
    length = random.randint(MIN_STR_LEN, MAX_STR_LEN)
    chars = string.ascii_letters + string.digits + "!@#$%^&*"
    return ''.join(random.choices(chars, k=length))


def encrypt_string(plaintext):
    """
    使用 iPN_StringEnc 加密字符串
    返回加密后的 bytes，失败返回 None
    """
    try:
        status, encrypted_data = iPN_StringEnc(ENCRYPTION_KEY, plaintext, ENCRYPTION_MODE)
        if status != 0 or encrypted_data is None:
            return None

        if isinstance(encrypted_data, bytes):
            return encrypted_data
        elif isinstance(encrypted_data, str):
            try:
                return bytes.fromhex(encrypted_data)
            except ValueError:
                try:
                    import base64
                    return base64.b64decode(encrypted_data, validate=True)
                except Exception:
                    return None
        else:
            return None
    except Exception as e:
        return None


def create_request(encrypted_bytes):
    """构建 UDP 请求数据包：Header + Encrypted Payload"""
    msg_len = len(encrypted_bytes)
    msg_type = get_random_msg_type()  # ✅ 每次随机选择 MSG_TYPE
    header = struct.pack(MSG_HEADER_FORMAT, START_CODE, msg_type, msg_len)
    return header + encrypted_bytes


def send_request(_):
    """
    发送单个加密请求
    参数 _ 是占位符
    """
    plaintext = generate_random_string()
    encrypted_data = encrypt_string(plaintext)

    if encrypted_data is None:
        return False  # 加密失败视为发送失败

    request_data = create_request(encrypted_data)
    try:
        with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock:
            sock.settimeout(TIMEOUT)
            sock.sendto(request_data, (SERVER_IP, SERVER_PORT))
        return True
    except socket.timeout:
        return False
    except Exception:
        return False


# ==================== 主逻辑 ====================
def send_batch():
    """发送一批加密请求，并统计 MSG_TYPE 分布"""
    total = REQUESTS_PER_BATCH
    successful = 0
    type_count = {1: 0, 3: 0}  # 统计 MSG_TYPE 分布

    start_time = time.time()

    with ThreadPoolExecutor(max_workers=NUM_THREADS) as executor:
        futures = []
        for _ in range(total):
            msg_type = get_random_msg_type()
            type_count[msg_type] += 1
            futures.append(executor.submit(send_request, None))

        for future in as_completed(futures):
            if future.result():
                successful += 1

    elapsed = time.time() - start_time
    qps = total / elapsed if elapsed > 0 else 0
    success_rate = (successful / total) * 100

    # 打印本次批次结果（含 MSG_TYPE 分布）
    now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    print(f"[{now}] Sent {total} [Type1:{type_count[1]}, Type3:{type_count[3]}], "
          f"Success: {successful}, QPS: {qps:.2f}, Success Rate: {success_rate:.1f}%, "
          f"Time: {elapsed:.2f}s")

    return successful, total


def main():
    # 测试加密函数
    print("🔍 正在测试 iPN_StringEnc 加密函数...")
    test_plaintext = "TestEncrypt123"
    try:
        status, result = iPN_StringEnc(ENCRYPTION_KEY, test_plaintext, ENCRYPTION_MODE)
        print(f"🧪 测试明文: {test_plaintext}")
        print(f"   Status: {status}")
        print(f"   Result: {repr(result)}")
        if status == 0 and result is not None:
            print("✅ 加密函数测试通过！")
        else:
            print("❌ 加密失败！请检查密钥或 string_enc.py 实现")
            exit(1)
    except Exception as e:
        print(f"❌ 测试加密函数异常: {e}")
        exit(1)

    print(f"\n🚀 定时任务开始...")
    print(f"目标：每 {SEND_INTERVAL} 秒发送 {REQUESTS_PER_BATCH} 个加密随机字符串")
    print(f"字符串长度: {MIN_STR_LEN} ~ {MAX_STR_LEN}")
    print(f"MSG_TYPE: 随机选择 1 或 3")
    print(f"服务器: {SERVER_IP}:{SERVER_PORT}, 线程数: {NUM_THREADS}")
    print(f"将持续运行 {RUN_DURATION // 3600} 小时...")

    start_time = time.time()
    total_success = 0
    total_requests = 0

    while time.time() - start_time < RUN_DURATION:
        success, count = send_batch()
        total_success += success
        total_requests += count

        # 精确控制发送间隔
        next_send_time = (start_time + (int((time.time() - start_time) / SEND_INTERVAL) + 1) * SEND_INTERVAL)
        sleep_time = next_send_time - time.time()
        if sleep_time > 0:
            time.sleep(sleep_time)

    # 最终统计
    run_time = time.time() - start_time
    overall_qps = total_requests / run_time if run_time > 0 else 0
    overall_success_rate = (total_success / total_requests * 100) if total_requests > 0 else 0

    print("\n" + "=" * 60)
    print("✅ 任务完成！")
    print(f"总运行时间: {run_time // 3600:.0f}h {(run_time % 3600) // 60:.0f}m")
    print(f"发送间隔: {SEND_INTERVAL} 秒")
    print(f"每批请求数: {REQUESTS_PER_BATCH}")
    print(f"总批次: {total_requests // REQUESTS_PER_BATCH}")
    print(f"总请求数: {total_requests:,}")
    print(f"成功数: {total_success:,}")
    print(f"成功率: {overall_success_rate:.2f}%")
    print(f"平均 QPS: {overall_qps:.2f}")
    print("=" * 60)


if __name__ == "__main__":
    main()